#include <torch/torch.h>
#include <LibDL/data/MNISTDataIterator/MNISTDataIterator.h>

#ifndef SWIG
typedef std::unique_ptr<torch::data::StatelessDataLoader<torch::data::datasets::MapDataset<torch::data::datasets::MNIST, torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor>>>, torch::data::samplers::RandomSampler>, std::default_delete<torch::data::StatelessDataLoader<torch::data::datasets::MapDataset<torch::data::datasets::MNIST, torch::data::transforms::Stack<torch::data::Example<at::Tensor, at::Tensor>>>, torch::data::samplers::RandomSampler>>> mnist_ptr;
#endif

namespace data {
    class MNISTDataLoader {
    public:
        explicit MNISTDataLoader(bool train,const std::string& pos);


        MNISTDataIterator *iterator() const {
            return new MNISTDataIterator(core);
        }

    private:
        mnist_ptr core;
    };
}
